package com.smh.aiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * @author shiminghui
 * @date 2025/6/3 17:04
 * @description: TODO
 */
@Component
public class MyMultiQueryExpander {


    @Resource
    private ChatModel dashscopeChatModel;

    public List<Query> expand(Query query) {

        MultiQueryExpander multiQueryExpander = MultiQueryExpander.builder()
                .numberOfQueries(3) // 设置要生成的查询数量
                .chatClientBuilder(ChatClient.builder(dashscopeChatModel))
                .includeOriginal(true) // 是否包含原始查询
                .build();

        List<Query> queries = multiQueryExpander.expand(query);

        return queries;
    }

}
